"""月之暗面Kimi模型适配器实现。"""

from typing import AsyncGenerator, Dict, Any, List
import aiohttp
import json

from .base import BaseModelAdapter, ModelConfig, ModelResponse


class KimiAdapter(BaseModelAdapter):
    """Kimi (Moonshot) 大模型适配器。"""
    
    # 定价（每千tokens，人民币）
    PRICING = {
        "moonshot-v1-8k": {"prompt": 0.012, "completion": 0.012},
        "moonshot-v1-32k": {"prompt": 0.024, "completion": 0.024},
        "moonshot-v1-128k": {"prompt": 0.060, "completion": 0.060}
    }
    
    def _initialize_client(self) -> None:
        """初始化Kimi客户端。"""
        self.base_url = self.config.api_base or "https://api.moonshot.cn/v1"
    
    async def generate(
        self,
        messages: List[Dict[str, str]],
        **kwargs
    ) -> ModelResponse:
        """生成响应。
        
        Args:
            messages: 对话消息
            **kwargs: 额外参数
            
        Returns:
            模型响应
        """
        url = f"{self.base_url}/chat/completions"
        
        headers = {
            "Authorization": f"Bearer {self.config.api_key}",
            "Content-Type": "application/json"
        }
        
        # Kimi支持的特殊功能
        tools = kwargs.pop("tools", None)
        search = kwargs.pop("search", False)
        
        data = {
            "model": self.config.model_name,
            "messages": messages,
            "temperature": self.config.temperature,
            "max_tokens": self.config.max_tokens,
            "top_p": self.config.top_p,
            "frequency_penalty": self.config.frequency_penalty,
            "presence_penalty": self.config.presence_penalty,
            "stream": False
        }
        
        # 添加Kimi特有功能
        if tools:
            data["tools"] = tools
        if search:
            data["search"] = search
        
        data.update(kwargs)
        data.update(self.config.extra_params)
        
        async with aiohttp.ClientSession() as session:
            async with session.post(
                url,
                headers=headers,
                json=data,
                timeout=aiohttp.ClientTimeout(total=self.config.timeout)
            ) as response:
                result = await response.json()
                
                if "error" in result:
                    raise Exception(f"Kimi API错误: {result['error'].get('message', 'Unknown error')}")
                
                choice = result["choices"][0]
                usage = result.get("usage", {})
                
                return ModelResponse(
                    content=choice["message"]["content"],
                    model=result.get("model", self.config.model_name),
                    usage={
                        "prompt_tokens": usage.get("prompt_tokens", 0),
                        "completion_tokens": usage.get("completion_tokens", 0),
                        "total_tokens": usage.get("total_tokens", 0)
                    },
                    finish_reason=choice.get("finish_reason", "stop"),
                    metadata={
                        "id": result.get("id", ""),
                        "created": result.get("created", 0)
                    }
                )
    
    async def stream(
        self,
        messages: List[Dict[str, str]],
        **kwargs
    ) -> AsyncGenerator[str, None]:
        """流式生成响应。
        
        Args:
            messages: 对话消息
            **kwargs: 额外参数
            
        Yields:
            内容块
        """
        url = f"{self.base_url}/chat/completions"
        
        headers = {
            "Authorization": f"Bearer {self.config.api_key}",
            "Content-Type": "application/json"
        }
        
        data = {
            "model": self.config.model_name,
            "messages": messages,
            "temperature": self.config.temperature,
            "max_tokens": self.config.max_tokens,
            "top_p": self.config.top_p,
            "stream": True
        }
        
        data.update(kwargs)
        data.update(self.config.extra_params)
        
        async with aiohttp.ClientSession() as session:
            async with session.post(
                url,
                headers=headers,
                json=data,
                timeout=aiohttp.ClientTimeout(total=self.config.timeout)
            ) as response:
                async for line in response.content:
                    if line:
                        line_str = line.decode('utf-8').strip()
                        if line_str.startswith("data: "):
                            data_str = line_str[6:]
                            if data_str == "[DONE]":
                                break
                            
                            try:
                                chunk = json.loads(data_str)
                                if "choices" in chunk and chunk["choices"]:
                                    delta = chunk["choices"][0].get("delta", {})
                                    content = delta.get("content", "")
                                    if content:
                                        yield content
                            except json.JSONDecodeError:
                                continue
    
    def count_tokens(self, text: str) -> int:
        """计算文本的token数量。
        
        Kimi对长文本处理优化，使用更精确的估算。
        
        Args:
            text: 输入文本
            
        Returns:
            Token数量
        """
        # 中文：约0.7个字符1个token（Kimi优化）
        chinese_chars = sum(1 for char in text if '\u4e00' <= char <= '\u9fff')
        # 英文：约4个字符1个token
        english_chars = len([c for c in text if c.isascii() and c.isalpha()])
        # 其他字符
        other_chars = len(text) - chinese_chars - english_chars
        
        return int(chinese_chars * 0.7 + english_chars / 4 + other_chars * 0.5)
    
    def estimate_cost(self, prompt_tokens: int, completion_tokens: int) -> float:
        """估算使用成本。
        
        Args:
            prompt_tokens: 提示词token数
            completion_tokens: 生成token数
            
        Returns:
            预估成本（美元）
        """
        # 根据token数自动选择模型
        total_tokens = prompt_tokens + completion_tokens
        
        if total_tokens <= 8000:
            model_key = "moonshot-v1-8k"
        elif total_tokens <= 32000:
            model_key = "moonshot-v1-32k"
        else:
            model_key = "moonshot-v1-128k"
        
        pricing = self.PRICING[model_key]
        prompt_cost = (prompt_tokens / 1000) * pricing["prompt"]
        completion_cost = (completion_tokens / 1000) * pricing["completion"]
        
        # 转换为美元
        cny_to_usd = 0.14
        return (prompt_cost + completion_cost) * cny_to_usd
    
    async def health_check(self) -> bool:
        """健康检查。
        
        Returns:
            是否健康
        """
        try:
            # 使用模型列表接口进行健康检查
            url = f"{self.base_url}/models"
            headers = {
                "Authorization": f"Bearer {self.config.api_key}"
            }
            
            async with aiohttp.ClientSession() as session:
                async with session.get(
                    url,
                    headers=headers,
                    timeout=aiohttp.ClientTimeout(total=5)
                ) as response:
                    if response.status == 200:
                        data = await response.json()
                        return "data" in data
                    return False
        except Exception:
            return False